package com.my.collection.tree.v2;

import java.util.Comparator;

/**
 * Created by ZhongGang
 * at 2018/10/26 10:10
 * <p>
 * 二叉搜索树 左子结点小于父结点，右子节点大于等于父结点
 */
public class BinarySearchTree<E> {
    private Node root;
    private Comparator<E> comparator;

    public BinarySearchTree(E root) {
        this(root, new Comparator<E>() {
            @Override
            public int compare(E o1, E o2) {
                Comparable<E> c1 = (Comparable<E>) o1;
                return c1.compareTo(o2);
            }
        });
    }

    public BinarySearchTree(E root, Comparator<E> comparator) {
        this.root = new Node(root, null, null);
        this.comparator = comparator;
    }

    public void add(E e) {
        if (e == null) {
            throw new NullPointerException();
        }

        Node node = new Node(e, null, null);
        Node tmp = this.root;
        while (tmp != null) {
            if (node.compareTo(tmp) < 0) {
                if (tmp.left != null) {
                    tmp = tmp.left;
                } else {
                    tmp.left = node;
                    return;
                }
            } else {
                if (tmp.right != null) {
                    tmp = tmp.right;
                } else {
                    tmp.right = node;
                    return;
                }
            }
        }

        this.root = node;
    }

    public void remove(E e) {
        if (e == null) {
            throw new NullPointerException();
        }

        Node parent = null;
        Node tmp = this.root;
        boolean left = false;
        while (tmp != null) {
            if (this.comparator.compare(tmp.data, e) < 0) {
                if (tmp.right != null) {
                    left = false;
                    parent = tmp;
                    tmp = tmp.right;
                } else {
                    return;
                }
            } else if (this.comparator.compare(tmp.data, e) > 0) {
                if (tmp.left != null) {
                    left = true;
                    parent = tmp;
                    tmp = tmp.left;
                } else {
                    return;
                }
            } else {
                if (tmp.left == null && tmp.right == null) {
                    //叶子节点
                    if (parent == null) {
                        this.root = null;
                    } else {
                        if (left) {
                            parent.left = null;
                        } else {
                            parent.right = null;
                        }
                    }
                }

                if (tmp.left != null && tmp.right == null) {
                    //只有左子树，没有右子树
                    if (parent == null) {
                        this.root = tmp.left;
                    } else {
                        if (left) {
                            parent.left = tmp.left;
                        } else {
                            parent.right = tmp.left;
                        }
                    }
                }

                if (tmp.left == null && tmp.right != null) {
                    //只有右子树，没有左子树
                    if (parent == null) {
                        this.root = tmp.right;
                    } else {
                        if (left) {
                            parent.left = tmp.right;
                        } else {
                            parent.right = tmp.right;
                        }
                    }
                }

                if (tmp.left != null && tmp.right != null) {
                    //左右子树都存在
                    Node t = tmp.right;
                    Node p = t;
                    while (t.left != null) {
                        p = t;
                        t = t.left;
                    }

                    if (parent == null) {
                        if (p == t) {
                            t.left = tmp.left;
                            this.root = t;
                        } else {
                            t.left = tmp.left;
                            p.left = t.right;
                            t.right = p;
                            this.root = t;
                        }
                    } else {
                        if (left) {
                            if (p == t) {
                                t.left = tmp.left;
                                parent.left = t;
                            } else {
                                t.left = tmp.left;
                                p.left = t.right;
                                t.right = p;
                                parent.left = t;
                            }
                        } else {
                            if (p == t) {
                                t.left = tmp.left;
                                parent.right = t;
                            } else {
                                t.left = tmp.left;
                                p.left = t.right;
                                t.right = p;
                                parent.right = t;
                            }
                        }
                    }
                }

                return;
            }
        }
    }

    public boolean contain(E e) {
        if (e == null) {
            throw new NullPointerException();
        }

        Node tmp = this.root;
        while (tmp != null) {
            if (this.comparator.compare(tmp.data, e) < 0) {
                if (tmp.right != null) {
                    tmp = tmp.right;
                } else {
                    return false;
                }
            } else if (this.comparator.compare(tmp.data, e) > 0) {
                if (tmp.left != null) {
                    tmp = tmp.left;
                } else {
                    return false;
                }
            } else {
                return true;
            }
        }

        return false;
    }

    private class Node implements Comparable<Node> {
        E data;
        Node left;
        Node right;

        public Node(E data, Node left, Node right) {
            this.data = data;
            this.left = left;
            this.right = right;
        }

        @Override
        public int compareTo(Node o) {
            return BinarySearchTree.this.comparator.compare(this.data, o.data);
        }
    }

}
